from fastapi import FastAPI, HTTPException, Depends, Query
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import List, Dict, Optional, Tuple
import threading
import datasets
import time
from contextlib import asynccontextmanager
from loguru import logger
import sys
import os
from huggingface_hub import login
from pathlib import Path
import numpy as np
import soundfile as sf
from datetime import datetime
from database import (
    SessionLocal, add_voice_data, get_voice_data_by_language,
    get_voice_data_count, get_db, get_voice_data_by_id,
    check_voice_data_exists
)

# 配置日志
logger.remove()  # 移除默认的处理器
logger.add(
    "logs/app.log",
    rotation="500 MB",    # 每个日志文件最大 500MB
    retention="10 days",  # 保留 10 天的日志
    compression="zip",    # 压缩旧的日志文件
    level="INFO",        # 日志级别
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>"
)
logger.add(sys.stdout, level="INFO")

BATCH_SIZE = 10  # 每批处理的数据量

# Hugging Face 配置
HF_ENDPOINT = os.getenv("HF_ENDPOINT", "https://huggingface.co")  # 默认使用官方站点
HF_MIRROR = os.getenv("HF_MIRROR", "")  # 镜像站点，例如 "https://hf-mirror.com"
OFFLINE_MODE = os.getenv("OFFLINE_MODE", "").lower() in ("true", "1", "yes")  # 离线模式

# 创建音频文件存储目录
AUDIO_DIR = Path("audio_files")
AUDIO_DIR.mkdir(exist_ok=True)

# 支持的语言列表, 选择了一些主要语言, 如果要支持所有语言可以从 README.md 中的支持的语言列表中选择
# 测试保留 2 个
# SUPPORTED_LANGUAGES = ['en', 'zh-CN']
# 常用语言
# SUPPORTED_LANGUAGES = ['en', 'zh-CN', 'fr', 'de', 'ja', 'ko', 'es', 'ru', 'it', 'pt', 'nl', 'pl', 'tr', 'ar', 'hi', 'vi', 'id', 'uk', 'th']、
# 所有语言
SUPPORTED_LANGUAGES = ['ab', 'af', 'am', 'ar', 'as', 'ast', 'az', 'ba', 'bas', 'be', 'bg', 'bn', 'br', 'ca', 'ckb', 'cnh', 'cs', 'cv', 'cy', 'da', 'de', 'dv', 'dyu', 'el', 'en', 'eo', 'es', 'et', 'eu', 'fa', 'fi', 'fr', 'fy-NL', 'ga-IE', 'gl', 'gn', 'ha', 'he', 'hi', 'hsb', 'ht', 'hu', 'hy-AM', 'ia', 'id', 'ig', 'is', 'it', 'ja', 'ka', 'kab', 'kk', 'kmr', 'ko', 'ky', 'lg', 'lij', 'lo', 'lt', 'ltg', 'lv', 'mdf', 'mhr', 'mk', 'ml', 'mn', 'mr', 'mrj', 'mt', 'myv', 'nan-tw', 'ne-NP', 'nhi', 'nl', 'nn-NO', 'nso', 'oc', 'or', 'os', 'pa-IN', 'pl', 'ps', 'pt', 'quy', 'rm-sursilv', 'rm-vallader', 'ro', 'ru', 'rw', 'sah', 'sat', 'sc', 'sk', 'skr', 'sl', 'sq', 'sr', 'sv-SE', 'sw', 'ta', 'te', 'th', 'ti', 'tig', 'tk', 'tok', 'tr', 'tt', 'tw', 'ug', 'uk', 'ur', 'uz', 'vi', 'vot', 'yi', 'yo', 'yue', 'zgh', 'zh-CN']

# 注意: zh-HK zh-TW zu zza 获取有问题

# 采集数据集的语言，之前是和 SUPPORTED_LANGUAGES 一致, 但是后面要求优先采集一些语言，所以就切割开来
COLLECT_LANGUAGES = SUPPORTED_LANGUAGES
# 只采集: 中亚语言
# COLLECT_LANGUAGES = ['uz', 'ky', 'ur', 'ps']

def save_audio(audio_data: dict, language: str, path: str) -> Optional[str]:
    """将音频数据保存为 WAV 文件"""
    try:
        # 获取音频数据
        array = np.array(audio_data['array'])
        sampling_rate = audio_data['sampling_rate']

        # 创建语言目录
        lang_dir = AUDIO_DIR / language
        lang_dir.mkdir(exist_ok=True)

        # 从path中提取文件名
        filename = Path(path).name
        filename_without_ext = Path(filename).stem
        wav_path = lang_dir / f"{filename_without_ext}.wav"

        # 保存为WAV文件
        logger.debug(f"Saving WAV file to {wav_path}")
        sf.write(str(wav_path), array, sampling_rate)

        logger.info(f"已保存音频: {wav_path}")
        return f"/audio/{language}/{filename_without_ext}.wav"
    except Exception as e:
        logger.error(f"Error saving audio: {str(e)}", exc_info=True)
        return None

class VoiceData(BaseModel):
    """语音数据模型"""
    id: int
    voice_path: str
    text: str
    language: str
    created_at: datetime

    class Config:
        from_attributes = True

    @classmethod
    def from_db(cls, db_item):
        """从数据库模型创建 Pydantic 模型"""
        return cls(
            id=db_item.id,
            voice_path=db_item.voice_path,
            text=db_item.text,
            language=db_item.language,
            created_at=db_item.created_at
        )

class VoiceDataPage(BaseModel):
    """分页返回的语音数据"""
    items: List[VoiceData]
    total: int
    total_pages: int

    class Config:
        from_attributes = True

    @classmethod
    def from_db_result(cls, result):
        """从数据库结果创建分页响应"""
        return cls(
            items=[VoiceData.from_db(item) for item in result["items"]],
            total=result["total"],
            total_pages=result["total_pages"]
        )

def validate_audio_data(audio_data: dict) -> Tuple[bool, str]:
    """验证音频数据的有效性"""
    try:
        # 检查必要字段
        if 'array' not in audio_data or 'sampling_rate' not in audio_data:
            return False, "缺少必要的音频数据字段"

        # 检查采样率
        if not isinstance(audio_data['sampling_rate'], (int, float)) or audio_data['sampling_rate'] <= 0:
            return False, f"无效的采样率: {audio_data['sampling_rate']}"

        # 检查音频数组
        audio_array = np.array(audio_data['array'])
        if audio_array.size == 0:
            return False, "空的音频数组"

        if not np.isfinite(audio_array).all():
            return False, "音频数组包含无效值(inf/nan)"

        if audio_array.ndim != 1:
            return False, f"音频数组维度错误: {audio_array.ndim}, 应该是1"

        # 检查音频长度（至少100ms）
        min_samples = int(audio_data['sampling_rate'] * 0.1)
        if audio_array.size < min_samples:
            return False, f"音频太短: {audio_array.size / audio_data['sampling_rate']:.2f}秒"

        return True, ""

    except Exception as e:
        return False, f"验证音频数据时发生错误: {str(e)}"

def validate_text(text: str) -> Tuple[bool, str]:
    """验证文本内容的有效性"""
    try:
        # 检查是否为空
        if not text or not text.strip():
            return False, "空文本"

        # 检查长度（至少2个字符）
        if len(text.strip()) < 2:
            return False, f"文本太短: {text}"

        # 检查是否只包含空白字符
        if text.isspace():
            return False, "文本只包含空白字符"

        return True, ""

    except Exception as e:
        return False, f"验证文本时发生错误: {str(e)}"

def process_dataset_batch(dataset, language: str, db) -> None:
    """处理数据集的一个批次"""
    try:
        batch_data = []
        for item in dataset:
            if len(batch_data) >= BATCH_SIZE:
                break

            # 检查是否有必要的字段
            if not (item.get('audio') and item.get('sentence') and item.get('path')):
                logger.warning(f"数据缺少必要字段: {item.keys()}")
                continue

            # 验证文本
            text_valid, text_error = validate_text(item['sentence'])
            if not text_valid:
                logger.warning(f"无效的文本数据: {text_error}, language={language}, text={item['sentence']}")
                continue

            # 验证音频
            audio_valid, audio_error = validate_audio_data(item['audio'])
            if not audio_valid:
                logger.warning(f"无效的音频数据: {audio_error}, language={language}, path={item['path']}")
                continue

            # 检查是否重复
            if check_voice_data_exists(db, item['sentence'], language):
                logger.debug(f"跳过重复数据: language={language}, text={item['sentence']}")
                continue

            batch_data.append(item)

        # 处理这批数据
        for item in batch_data:
            # 保存音频文件
            voice_path = save_audio(item['audio'], language, item['path'])
            if voice_path:
                # 添加到数据库
                add_voice_data(db, voice_path, item['sentence'], language)
                logger.info(f"已保存数据: language={language}, text={item['sentence']}")

    except Exception as e:
        logger.error(f"Error processing batch for language {language}: {str(e)}")


@asynccontextmanager
async def lifespan(app: FastAPI):
    """启动时开始获取数据，关闭时清理资源"""

    # 初始化 Hugging Face 认证
    if not init_huggingface():
        logger.error("Failed to initialize Hugging Face authentication. The application might not work correctly.")

    logger.info("Starting data fetching thread...")
    # 启动后台数据获取线程
    thread = threading.Thread(target=fetch_data, daemon=True)
    thread.start()
    yield
    logger.info("Shutting down application...")


app = FastAPI(title="Voice Dataset API", lifespan=lifespan)

# 挂载静态文件目录
app.mount("/audio", StaticFiles(directory=str(AUDIO_DIR)), name="audio")

current_language_index = 0  # 当前处理的语言索引

def init_huggingface():
    """
    初始化 Hugging Face 认证

    1. 访问 https://huggingface.co/settings/tokens
    2. 如果没有账号，需要先注册
    3. 登录后创建一个新的 token (建议选择"read"权限即可)
    4. 复制生成的token
    5. 设置环境变量：
       - HF_TOKEN: 访问令牌（在线模式必需）
       - HF_ENDPOINT: [可选] API端点, 默认 https://huggingface.co
       - HF_MIRROR: [可选] 镜像站点
       - OFFLINE_MODE: [可选] 离线模式，设置为 true 则不进行在线验证
           需要注意的是：
           1. 在线模式需要 Hugging Face 的 token
           2. 离线模式可以直接使用镜像站点，无需 token
           3. 不同的镜像站点可能有不同的访问方式，可能需要根据具体的镜像站点调整代码

    PowerShell 配置:
    $env:HF_TOKEN = "你的token"                    # 在线模式必需
    $env:HF_ENDPOINT = "https://huggingface.co"    # 可选
    $env:HF_MIRROR = "https://hf-mirror.com"       # 可选
    $env:OFFLINE_MODE = "true"                     # 可选，设置为 true 则不进行在线验证

    Ubuntu 配置:
    export HF_TOKEN="你的token"                     # 在线模式必需
    export HF_ENDPOINT="https://huggingface.co"     # 可选
    export HF_MIRROR="https://hf-mirror.com"        # 可选
    export OFFLINE_MODE="true"                      # 可选，设置为 true 则不进行在线验证
    """
    # 如果是离线模式，跳过登录验证
    if OFFLINE_MODE:
        logger.info("Running in offline mode, skipping Hugging Face authentication")
        return True

    token = os.getenv("HF_TOKEN")
    if not token:
        logger.warning("HF_TOKEN environment variable not set. You need to set it for online mode.")
        logger.info("Set OFFLINE_MODE=true to skip authentication if you're using a mirror.")
        return False

    try:
        # 设置 API 端点
        if HF_MIRROR:
            logger.info(f"使用镜像站点: {HF_MIRROR}")
            os.environ["HF_ENDPOINT"] = HF_MIRROR
            datasets.config.HF_ENDPOINT = HF_MIRROR  # 设置 datasets 库的端点

        # 登录到 Hugging Face
        login(token=token, add_to_git_credential=False)
        logger.info("Successfully logged in to Hugging Face")
        return True
    except Exception as e:
        logger.error(f"Failed to login to Hugging Face: {str(e)}")
        logger.info("If you're using a mirror, consider setting OFFLINE_MODE=true to skip authentication")
        return False

def fetch_data():
    """持续获取数据的后台任务"""
    global current_language_index

    while True:
        try:
            # 获取当前要处理的语言
            language = COLLECT_LANGUAGES[current_language_index]
            logger.info(f"Processing language: {language}")

            # 创建新的数据库会话
            db = SessionLocal()
            try:
                # 加载数据集
                dataset = datasets.load_dataset(
                    "mozilla-foundation/common_voice_17_0",
                    language,
                    split="train",
                    streaming=True,
                    trust_remote_code=True
                )

                # 处理一批数据
                process_dataset_batch(dataset, language, db)
            finally:
                db.close()

            # 更新语言索引
            current_language_index = (current_language_index + 1) % len(COLLECT_LANGUAGES)

            # 休眠一段时间
            time.sleep(1)
        except Exception as e:
            logger.error(f"Error fetching data: {str(e)}")


@app.get("/counts")
def get_counts(db: SessionLocal = Depends(get_db)):
    """获取所有语种的数据条数"""
    counts = {}
    for language in SUPPORTED_LANGUAGES:
        counts[language] = get_voice_data_count(db, language)
    return counts

@app.get("/count/{language}")
def get_count(language: str, db: SessionLocal = Depends(get_db)):
    """获取指定语言的数据条数"""
    if language not in SUPPORTED_LANGUAGES:
        raise HTTPException(status_code=400, detail="Language not supported")
    return {"count": get_voice_data_count(db, language)}

@app.get("/data/{language}", response_model=VoiceDataPage)
def get_data(
    language: str,
    page: int = Query(default=1, ge=1, description="页码, 从1开始"),
    page_size: int = Query(default=20, ge=1, le=100, description="每页数据量, 最大100")
):
    """获取指定语言的数据，支持分页"""
    if language not in SUPPORTED_LANGUAGES:
        raise HTTPException(status_code=400, detail="Language not supported")

    db = SessionLocal()
    try:
        result = get_voice_data_by_language(db, language, page, page_size)
        return VoiceDataPage.from_db_result(result)
    finally:
        db.close()

@app.get("/data/id/{id}", response_model=VoiceData)
def get_data_by_id(id: int):
    """通过ID获取单条语音数据"""
    db = SessionLocal()
    try:
        data = get_voice_data_by_id(db, id)
        if not data:
            raise HTTPException(status_code=404, detail="Data not found")
        return VoiceData.from_db(data)
    finally:
        db.close()

if __name__ == "__main__":
    import uvicorn
    logger.info("Starting Voice Dataset API server...")
    uvicorn.run(app, host="0.0.0.0", port=9234)
